import sys
import os
# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import odeint
import pendulum_config

# 从pendulum_config模块获取参数
l = pendulum_config.l

def plot_phase_diagram(k, theta0, omega0, t_max=20):
    """
    绘制相图（角度-角速度图）
    """
    t = np.linspace(0, t_max, 1000)
    y0 = [theta0, omega0]
    sol = odeint(pendulum_config.pendulum_eq, y0, t, args=(k,))
    theta = sol[:, 0]
    omega = sol[:, 1]
    
    plt.figure(figsize=(10, 8))
    plt.plot(theta, omega, 'b-', linewidth=2)
    plt.xlabel('角度 θ (弧度)', fontsize=14)
    plt.ylabel('角速度 dθ/dt (弧度/秒)', fontsize=14)
    plt.title(f'有阻尼单摆相图 (k = {k:.3f} kg/s)', fontsize=16)
    plt.grid(True, alpha=0.3)
    plt.xlim(-2*np.pi, 2*np.pi)
    
    # 标记起始点
    plt.plot(theta0, omega0, 'ro', markersize=10, label='起始点')
    plt.legend()
    
    # 添加平衡点
    equilibrium_points = [0, np.pi, -np.pi]
    for point in equilibrium_points:
        plt.plot(point, 0, 'ks', markersize=8)
    
    plt.text(0.2, 0.2, '稳定平衡点', fontsize=12)
    plt.text(np.pi+0.2, 0.2, '不稳定平衡点', fontsize=12)
    plt.text(-np.pi+0.2, 0.2, '不稳定平衡点', fontsize=12)
    
    plt.tight_layout()
    # 在当前目录保存图片
    plt.savefig('damped_pendulum_phase.png', dpi=150, bbox_inches='tight')
    plt.show()